package com.zjp.tree;

import cn.hutool.core.collection.CollUtil;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public final class TreeUtil {

    public static <T, E> List<TreeNode<E>> build(List<T> list, Function<T, TreeNode<E>> data) {
        if (CollUtil.isEmpty(list)) {
            return Collections.emptyList();
        }
        Set<E> ids = new HashSet<>(list.size());
        List<TreeNode<E>> nodeList = new ArrayList<>(list.size());

        for (T t : list) {
            TreeNode<E> apply = data.apply(t);
            ids.add(apply.getId());
            nodeList.add(apply);
        }

        Map<E, List<TreeNode<E>>> map = nodeList.stream().collect(Collectors.groupingBy(TreeNode::getPId));
        // 顶级节点
        List<E> rootNode = new ArrayList<>();
        final String ROOT_NODE_ID = "0";

        for (TreeNode<E> node : nodeList) {
            E id = node.getId();
            E pId = node.getPId();
            // 没有上级
            if (!ids.contains(pId) || ROOT_NODE_ID.equals(pId.toString())) {
                rootNode.add(id);
            }
            // 设置子级
            node.setChildren(map.getOrDefault(node.getId(), Collections.emptyList()));
        }

        return nodeList.stream().filter(bean -> rootNode.contains(bean.getId())).collect(Collectors.toList());
    }

}
